#coding=utf-8
"""
!/usr/bin/python3
@Author: Gao Shuo
@Time: 2018/8/10 11:24 
@ReadMe:
    Input: 输入位置（可以逗号分割多个目录）、输出位置、图片宽度、图片高度
    Output: 各label的图片数、像素数、权重
"""

from PIL import Image
import os.path
import pandas as pd
import sys
import numpy as np


def img_in_path(path):
    li = os.listdir(path)
    res = []
    if len(li) == 0:
        return res
    else:
        for path_me in li:
            path_me_os =  os.path.join(path, path_me)
            res.append(path_me_os)
    return res

if __name__ == "__main__":

    path = sys.argv[1]
    dest = sys.argv[2]

    width = int(sys.argv[3])
    height = int(sys.argv[4])
    li_path = path.split(',')
    print li_path
    li_filepath = []
    for i in range(len(li_path)):
        li_this = img_in_path (li_path[i])
        li_filepath.extend(li_this)

    # 合并：检查颜色并计算
    dic_pixel = {}
    dic_pic = {}
    for filein in li_filepath:
        # filein = os.path.join(path, filepath)
        im = Image.open(filein)
        ima = np.array(im)
        # 本图有多少种颜色
        dic_this = {}
        for i in np.unique(ima):
            dic_this[i] = np.sum(ima==i)
        # for i in range(width):
        #     for j in range(height):
        #         color =  im.getpixel(xy = (i,j))
        #         if color in dic_this.keys():
        #             dic_this[color] += 1
        #         else:
        #             dic_this[color] =1
        for k,v in dic_this.items():
            if k in dic_pixel:
                dic_pixel[k] += v
            else:
                dic_pixel[k] = v
            if k in dic_pic:
                dic_pic[k] += 1
            else:
                dic_pic[k] = 1

        print filein, "  Color list:  ", sorted(dic_this.keys())
        print 'Finished: ', filein
    # 权重：c = median_freq / freq(c)
    df_pic = pd.DataFrame(dic_pic.items(), columns=['color', 'num_pic'])
    df_pixel = pd.DataFrame(dic_pixel.items(), columns=['color', 'num_pixel'])
    df1 = pd.merge(df_pic, df_pixel, on = 'color')
    df1["freq"] = ( df1["num_pixel"].astype(float) / (width * height * df1["num_pic"]))
    mean_freq = np.mean(df1["freq"])
    df1["class_weighting"] =  mean_freq  / df1["freq"]
    df1["class_weighting_txt"] =  "    class_weighting: " + str(df1["class_weighting"])

    print df1["freq"]
    print df1["class_weighting"]
    df1.to_csv(dest)

